/*
 ref link:
    https://blog.csdn.net/leixiaohua1020/article/details/39702113
    https://blog.csdn.net/xiaolong1126626497/article/details/104919095
    https://zhuanlan.zhihu.com/p/102489856
*/

#include <stdio.h>
#include <stdlib.h>

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
#include <libavdevice/avdevice.h>

#define YUV_WIDTH 352
#define YUV_HEIGHT 288


int main(int argc, char* argv[]){
    if(argc < 2){
        printf("Please input YUV file!\n");
        return -1;
    }

    int ret = 0;
    FILE* fp = NULL;
    char* url = NULL;

    AVFormatContext *pfmt_ctx = NULL;
    AVInputFormat *infmt = NULL;
    AVCodecParameters *pcodec_par = NULL;
    AVCodecContext *pcodec_ctx = NULL;
    AVCodec *pcodec = NULL;
    AVPacket *ppkt = NULL;
    AVFrame *pfrm = NULL;
    AVFrame *pyuv_frm = NULL;

    struct SwsContext *psws_ctx = NULL;
    int video_idx = -1;

    url = argv[1];

    fp = fopen(url, "wb");
    if(fp == NULL){
        fprintf(stderr, "File open ERROR\n");
        return -1;
    }

    int frame_idx = 0;

    avdevice_register_all();
    
    if ((pfmt_ctx = avformat_alloc_context()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((infmt = av_find_input_format("video4linux2")) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_find_input_format ERROR \n");
        return AVERROR(EINVAL);
    }

    if((ret = avformat_open_input(&pfmt_ctx, "/dev/video0", infmt, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_open_input ERROR \n");
        return AVERROR(EINVAL);
    }

    if ((ret = avformat_find_stream_info(pfmt_ctx, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_find_stream_info ERROR \n");
        return AVERROR(EINVAL);
    }

    for (int i=0; i<pfmt_ctx->nb_streams; i++) {
        if (pfmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_idx = i;
            break;
        }
    }

    if (video_idx == -1) {
        fprintf(stderr, "Can not find video index \n");
        return -1;
    }

    pcodec_par = pfmt_ctx->streams[video_idx]->codecpar;

    if ((pcodec = avcodec_find_decoder(pcodec_par->codec_id)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder ERROR \n");
        return AVERROR(ret);
    }
    
    if ((pcodec_ctx = avcodec_alloc_context3(pcodec)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((ret = avcodec_parameters_to_context(pcodec_ctx, pcodec_par)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_to_context ERROR \n");
        return AVERROR(ret);
    }
    
    pcodec_ctx->time_base.den = 25;
    pcodec_ctx->time_base.num = 1;

    if ((ret = avcodec_open2(pcodec_ctx, pcodec, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR \n");
        return AVERROR(ret);
    }

    if ((ppkt = av_packet_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((pfrm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((pyuv_frm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((fp = fopen(url, "wb")) == NULL) {
        fprintf(stderr, "Can not open file %s \n", url);
        return -2;
    }

    if ((psws_ctx = sws_getContext(pcodec_ctx->width, pcodec_ctx->height, pcodec_ctx->pix_fmt, YUV_WIDTH, YUV_HEIGHT, AV_PIX_FMT_YUV420P, SWS_FAST_BILINEAR, NULL, NULL, NULL)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "sws_getContext ERROR \n");
        return AVERROR(EINVAL);
    }

    if ((ret = av_image_alloc(pyuv_frm->data, pyuv_frm->linesize, YUV_WIDTH, YUV_HEIGHT, AV_PIX_FMT_YUV420P, 1)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_image_alloc ERROR \n");
        return AVERROR(EINVAL);
    }

    while ((ret = av_read_frame(pfmt_ctx, ppkt)) >= 0) {
        if (ppkt->stream_index == video_idx) {
            if ((ret = avcodec_send_packet(pcodec_ctx, ppkt)) < 0) {
                av_log(NULL, AV_LOG_ERROR, "avcodec_send_packet ERROR %s \n", av_err2str(ret));
                return AVERROR(ret);
            }

            while (ret >= 0) {
                ret = avcodec_receive_frame(pcodec_ctx, pfrm);
                if (ret == AVERROR_EOF || ret == AVERROR(EAGAIN)) {
                    break;
                } else if (ret < 0) {
                    av_log(NULL, AV_LOG_ERROR, "avcodec_receive_frame ERROR \n");
                    return AVERROR(ret);
                }

                sws_scale(psws_ctx, (const uint8_t * const *) pfrm->data, pfrm->linesize, 0, pcodec_ctx->height, pyuv_frm->data, pyuv_frm->linesize);

                fwrite(pyuv_frm->data[0], 1, YUV_WIDTH * YUV_HEIGHT, fp);
                fwrite(pyuv_frm->data[1], 1, YUV_WIDTH * YUV_HEIGHT / 4, fp);
                fwrite(pyuv_frm->data[2], 1, YUV_WIDTH * YUV_HEIGHT / 4, fp);

                frame_idx++;
                fprintf(stdout, "Current frame index: %d \n", frame_idx);
                av_frame_unref(pfrm);

            }
        }
        av_packet_unref(ppkt);
        if (frame_idx == 100) {
            break;
        }
    }

    fclose(fp);
    avformat_close_input(&pfmt_ctx);
    avformat_free_context(pfmt_ctx);
    av_frame_free(&pfrm);
    av_frame_free(&pyuv_frm);
    av_packet_free(&ppkt);
    sws_freeContext(psws_ctx);

    return 0;
}
